import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


# ---------------------- SuperPoint 模型定义 ----------------------
class SuperPoint(nn.Module):

    def __init__(self, config):
        super().__init__()
        self.nms_radius = config['nms_radius']
        self.max_keypoints = config['max_keypoints']
        self.descriptor_dim = config['descriptor_dim']
        self.relu = nn.ReLU(inplace=True)

        # 特征提取层
        self.conv1a = nn.Conv2d(1, 64, 3, stride=1, padding=1)
        self.conv1b = nn.Conv2d(64, 64, 3, stride=2, padding=1)
        self.conv2a = nn.Conv2d(64, 128, 3, stride=1, padding=1)
        self.conv2b = nn.Conv2d(128, 128, 3, stride=2, padding=1)
        self.conv3a = nn.Conv2d(128, 256, 3, stride=1, padding=1)
        self.conv3b = nn.Conv2d(256, 256, 3, stride=2, padding=1)
        self.convPa = nn.Conv2d(256, 65, 1, stride=1, padding=0)
        self.convPb = nn.Conv2d(256, self.descriptor_dim, 1, stride=1, padding=0)

        # 加载预训练权重
        self.load_state_dict(torch.load('superpoint_v1.pth', map_location='cpu'))
        self.eval()

    def forward(self, x):
        x = self.relu(self.conv1a(x))
        x = self.relu(self.conv1b(x))
        x = self.relu(self.conv2a(x))
        x = self.relu(self.conv2b(x))
        x = self.relu(self.conv3a(x))
        x = self.relu(self.conv3b(x))
        scores = self.convPa(x)
        descriptors = self.convPb(x)
        return self._process_output(scores, descriptors)

    def _process_output(self, scores, descriptors):
        scores = scores.permute(0, 2, 3, 1).squeeze()[:, :, :65]
        descriptors = descriptors.permute(0, 2, 3, 1).squeeze()
        return {'scores': scores, 'descriptors': descriptors}
